# This is free and unencumbered software released into the public domain.
#
# Anyone is free to copy, modify, publish, use, compile, sell, or distribute
# this software, either in source code form or as a compiled binary, for any
# purpose, commercial or non-commercial, and by any means.
#
# In jurisdictions that recognize copyright laws, the author or authors of this
# software dedicate any and all copyright interest in the software to the public
# domain. We make this dedication for the benefit of the public at large and to
# the detriment of our heirs and successors. We intend this dedication to be an
# overt act of relinquishment in perpetuity of all present and future rights to
# this software under copyright law.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import torch
from min_dalle import MinDalle
from PIL import Image
import numpy as np

class MinDalleNode:
    # Model cache
    _model = None
    _model_is_mega = False

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "An astronaut riding a horse"}),
                "seed": ("INT", {"default": 42, "min": 0, "max": 18446744073709551615}),
                "top_k": ("INT", {"default": 256}),
                "supercondition_factor": ("INT", {"default": 16}),
                "mega": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("PIL_IMAGE",)
    FUNCTION = "generate"

    CATEGORY = "min-dalle"

    def generate(self, prompt, seed, top_k, supercondition_factor, mega):
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

        # Cache the model
        if self._model is None or self._model_is_mega != mega:
            self._model = MinDalle(
                dtype=torch.float32,
                device=device,
                is_mega=mega,
                is_reusable=True
            )
            self._model_is_mega = mega

        # Generate the image
        image = self._model.generate_image(
            text=prompt,
            seed=seed,
            grid_size=1,
            top_k=top_k,
            supercondition_factor=supercondition_factor,
            is_seamless=False,
            temperature=1.0,
            is_verbose=False
        )

        return ([image],)
